package dataframe

import (
	"fmt"
	"hash/fnv"

	"gitee.com/lilihli/gota/series"
)

type joinType int

const (
	innerJoin joinType = iota
	leftJoin
	rightJoin
	outerJoin
)

type joinContext struct {
	left, right      *DataFrame
	iKeysL, iKeysR   []int
	nonKeyL, nonKeyR []int
	keySetL, keySetR map[int]bool
	resultCols       []series.Series
}

func (df DataFrame) InnerJoin(b DataFrame, keys ...string) DataFrame {
	return df.join(b, innerJoin, keys)
}

func (df DataFrame) LeftJoin(b DataFrame, keys ...string) DataFrame {
	return df.join(b, leftJoin, keys)
}

func (df DataFrame) RightJoin(b DataFrame, keys ...string) DataFrame {
	return df.join(b, rightJoin, keys)
}

func (df DataFrame) OuterJoin(b DataFrame, keys ...string) DataFrame {
	return df.join(b, outerJoin, keys)
}

func (df DataFrame) join(other DataFrame, typ joinType, keys []string) DataFrame {
	if len(keys) == 0 {
		return DataFrame{Err: fmt.Errorf("join keys not specified")}
	}

	ctx := &joinContext{
		left:    &df,
		right:   &other,
		keySetL: make(map[int]bool),
		keySetR: make(map[int]bool),
	}

	if err := ctx.initKeys(keys); err != nil {
		return DataFrame{Err: err}
	}

	switch typ {
	case innerJoin:
		// 内连接：智能选择小表构建哈希表
		buildDF, buildKeys, probeDF, probeKeys, isRightJoin := ctx.selectTablesForInner()
		hash := ctx.buildHashTable(buildDF, buildKeys)
		ctx.resultCols = ctx.prepareResultColumns()
		ctx.executeInner(probeDF, probeKeys, isRightJoin, hash)

	case leftJoin:
		// 左连接：右表构建哈希表，左表探测
		hash := ctx.buildHashTable(ctx.right, ctx.iKeysR)
		ctx.resultCols = ctx.prepareResultColumns()
		ctx.executeLeft(ctx.left, hash)

	case rightJoin:
		// 右连接：左表构建哈希表，右表探测
		hash := ctx.buildHashTable(ctx.left, ctx.iKeysL)
		ctx.resultCols = ctx.prepareResultColumnsForRight()
		ctx.executeRight(hash)

	case outerJoin:
		// 外连接：右表构建哈希表，左表探测
		hash := ctx.buildHashTable(ctx.right, ctx.iKeysR)
		ctx.resultCols = ctx.prepareResultColumns()
		ctx.executeOuter(ctx.left, ctx.iKeysL, hash)
	}

	return New(ctx.resultCols...)
}

func (ctx *joinContext) initKeys(keys []string) error {
	ctx.iKeysL = make([]int, len(keys))
	ctx.iKeysR = make([]int, len(keys))
	for i, key := range keys {
		if idx := ctx.left.colIndex(key); idx >= 0 {
			ctx.iKeysL[i] = idx
			ctx.keySetL[idx] = true
		} else {
			return fmt.Errorf("key %q not found in left DataFrame", key)
		}
		if idx := ctx.right.colIndex(key); idx >= 0 {
			ctx.iKeysR[i] = idx
			ctx.keySetR[idx] = true
		} else {
			return fmt.Errorf("key %q not found in right DataFrame", key)
		}
	}

	for i := 0; i < ctx.left.ncols; i++ {
		if !ctx.keySetL[i] {
			ctx.nonKeyL = append(ctx.nonKeyL, i)
		}
	}
	for i := 0; i < ctx.right.ncols; i++ {
		if !ctx.keySetR[i] {
			ctx.nonKeyR = append(ctx.nonKeyR, i)
		}
	}
	return nil
}

// selectTablesForInner 为内连接选择小表作为构建表
func (ctx *joinContext) selectTablesForInner() (build *DataFrame, buildKeys []int, probe *DataFrame, probeKeys []int, isRightJoin bool) {
	if ctx.left.nrows <= ctx.right.nrows {
		return ctx.left, ctx.iKeysL, ctx.right, ctx.iKeysR, false
	}
	return ctx.right, ctx.iKeysR, ctx.left, ctx.iKeysL, true
}

func (ctx *joinContext) buildHashTable(df *DataFrame, keyIdx []int) map[uint64][]int {
	table := make(map[uint64][]int, df.nrows)
	for i := 0; i < df.nrows; i++ {
		key := ctx.hashKey(df, keyIdx, i)
		table[key] = append(table[key], i)
	}
	return table
}

// 核心修复：executeInner正确处理小表构建的方向
func (ctx *joinContext) executeInner(probe *DataFrame, probeKeys []int, isRightJoin bool, hash map[uint64][]int) {
	for i := 0; i < probe.nrows; i++ {
		key := ctx.hashKey(probe, probeKeys, i)
		if matches, ok := hash[key]; ok {
			for _, j := range matches {
				if isRightJoin {
					// probe是左表，build是右表：i=左行, j=右行
					ctx.appendMatchedRow(i, j)
				} else {
					// probe是右表，build是左表：j=左行, i=右行
					ctx.appendMatchedRow(j, i)
				}
			}
		}
	}
}

// executeLeft 左连接：保留所有左表行
func (ctx *joinContext) executeLeft(probe *DataFrame, hash map[uint64][]int) {
	for i := 0; i < probe.nrows; i++ {
		key := ctx.hashKey(probe, ctx.iKeysL, i)
		if matches, ok := hash[key]; ok {
			for _, j := range matches {
				ctx.appendMatchedRow(i, j)
			}
		} else {
			ctx.appendLeftUnmatchedRow(i)
		}
	}
}

// executeRight 右连接：保留所有右表行
func (ctx *joinContext) executeRight(hash map[uint64][]int) {
	for i := 0; i < ctx.right.nrows; i++ {
		key := ctx.hashKey(ctx.right, ctx.iKeysR, i)
		if matches, ok := hash[key]; ok {
			for _, j := range matches {
				ctx.appendMatchedRow(j, i)
			}
		} else {
			ctx.appendRightUnmatchedRow(i)
		}
	}
}

// 核心修复：executeOuter专门处理外连接，不移除isRightJoin参数混淆逻辑
func (ctx *joinContext) executeOuter(probe *DataFrame, probeKeys []int, hash map[uint64][]int) {
	matchedRight := make([]bool, ctx.right.nrows)

	for i := 0; i < probe.nrows; i++ {
		key := ctx.hashKey(probe, probeKeys, i)
		if matches, ok := hash[key]; ok {
			for _, j := range matches {
				// 外连接：始终是左行+右行
				ctx.appendMatchedRow(i, j)
				matchedRight[j] = true
			}
		} else {
			ctx.appendLeftUnmatchedRow(i)
		}
	}

	for i := 0; i < ctx.right.nrows; i++ {
		if !matchedRight[i] {
			ctx.appendRightUnmatchedRow(i)
		}
	}
}

func (ctx *joinContext) appendMatchedRow(leftRow, rightRow int) {
	idx := 0
	for _, colIdx := range ctx.iKeysL {
		elem, isNaN := ctx.left.columns[colIdx].Elem(leftRow)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
	for _, colIdx := range ctx.nonKeyL {
		elem, isNaN := ctx.left.columns[colIdx].Elem(leftRow)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
	for _, colIdx := range ctx.nonKeyR {
		elem, isNaN := ctx.right.columns[colIdx].Elem(rightRow)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
}

func (ctx *joinContext) appendLeftUnmatchedRow(row int) {
	idx := 0
	for _, colIdx := range ctx.iKeysL {
		elem, isNaN := ctx.left.columns[colIdx].Elem(row)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
	for _, colIdx := range ctx.nonKeyL {
		elem, isNaN := ctx.left.columns[colIdx].Elem(row)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
	for range ctx.nonKeyR {
		ctx.resultCols[idx].AppendOne(nil, true)
		idx++
	}
}

func (ctx *joinContext) appendRightUnmatchedRow(row int) {
	idx := 0
	for _, colIdx := range ctx.iKeysR {
		elem, isNaN := ctx.right.columns[colIdx].Elem(row)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
	for range ctx.nonKeyL {
		ctx.resultCols[idx].AppendOne(nil, true)
		idx++
	}
	for _, colIdx := range ctx.nonKeyR {
		elem, isNaN := ctx.right.columns[colIdx].Elem(row)
		ctx.resultCols[idx].AppendOne(elem, isNaN)
		idx++
	}
}

func (ctx *joinContext) prepareResultColumns() []series.Series {
	cols := make([]series.Series, 0, ctx.left.ncols+ctx.right.ncols-len(ctx.iKeysL))
	for _, idx := range ctx.iKeysL {
		cols = append(cols, ctx.left.columns[idx].Empty())
	}
	for _, idx := range ctx.nonKeyL {
		cols = append(cols, ctx.left.columns[idx].Empty())
	}
	for _, idx := range ctx.nonKeyR {
		cols = append(cols, ctx.right.columns[idx].Empty())
	}
	return cols
}

func (ctx *joinContext) prepareResultColumnsForRight() []series.Series {
	cols := make([]series.Series, 0, ctx.left.ncols+ctx.right.ncols-len(ctx.iKeysR))
	for _, idx := range ctx.iKeysR {
		cols = append(cols, ctx.right.columns[idx].Empty())
	}
	for _, idx := range ctx.nonKeyL {
		cols = append(cols, ctx.left.columns[idx].Empty())
	}
	for _, idx := range ctx.nonKeyR {
		cols = append(cols, ctx.right.columns[idx].Empty())
	}
	return cols
}

func (ctx *joinContext) hashKey(df *DataFrame, keyIdx []int, row int) uint64 {
	h := fnv.New64a()
	for _, i := range keyIdx {
		elem, isNaN := df.columns[i].Elem(row)
		if isNaN {
			h.Write([]byte{0})
		} else {
			fmt.Fprintf(h, "%v", elem)
			h.Write([]byte{0}) // 分隔符避免边界情况
		}
	}
	return h.Sum64()
}
